(Draft)[Main][feat] Support overlapping A2A Combine backprop with wgrad GEMM#3792
(Draft)[Main][feat] Support overlapping A2A Combine backprop with wgrad GEMM#3792Wohox wants to merge 7 commits intoNVIDIA:mainfrom
Conversation
Wohox
changed the title
Wohox
changed the title
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ReviewThe approach is solid — separating expert wgrad onto a dedicated CUDA stream to overlap with the dispatch backward (A2A combine backprop) is well-structured with the two autograd function design. One inline suggestion on Missing test coverageThere are no tests for the new
|
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
|
/claude review |
There was a problem hiding this comment.
Light review — overall this looks well-structured. The CUDA stream synchronization pattern is correct and the FSDP integration is clean.
One observation on test parametrization: test_delay_wgrad_compute_for_te_grouped_gemm and the FSDP variant parametrize flex_dispatcher_backend independently of dispatcher_type, which creates redundant test combinations (e.g., dispatcher_type="alltoall" × all flex backends) where the backend parameter is unused. This multiplies CI cost without adding coverage. Consider either collapsing these into a single parametrize that yields valid (dispatcher_type, backend) pairs, or using pytest.mark.parametrize with indirect to skip invalid combos.
megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
Outdated
Show resolved
Hide resolved
|
/claude review |
megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
Outdated
Show resolved
Hide resolved
|
Light review — one bug found, otherwise looks good. Bug: In Everything else looks solid — the CUDA stream synchronization logic is correct, the mutual exclusivity assertion between |
|
/claude review |
| wgrad_stream.wait_event(event) | ||
| with torch.cuda.stream(wgrad_stream): | ||
| with torch.cuda.nvtx.range("delayed_expert_wgrad"): | ||
| module.backward_dw(routed_experts=True, shared_experts=False) | ||
| event.record(wgrad_stream) | ||
|
|
||
| torch.cuda.current_stream().wait_event(event) |
There was a problem hiding this comment.
Bug: When moe_latent_size is configured, backward_dw() internally launches fc2_latent_proj.backward_dw() on comm_stream (line 585-586). However, the event here is only recorded on wgrad_stream, so torch.cuda.current_stream().wait_event(event) on line 668 does not wait for the comm_stream work to finish. This could lead to a data race where the main stream proceeds (e.g., starts the next iteration) before the latent projection weight gradients are fully computed.
You'd need to either:
- Also synchronize
comm_streamback towgrad_stream(or main stream) before recording the event, or - Add a validation assertion that
delay_wgrad_compute_for_te_grouped_gemmandmoe_latent_sizeare mutually exclusive (if that combination isn't intended to be supported yet).
|
/claude review |
There was a problem hiding this comment.
Looks good overall — clean separation of the delayed wgrad path into two autograd functions, and the FSDP integration is well-structured.
One potential bug flagged as an inline comment: when moe_latent_size is configured, backward_dw() dispatches work to comm_stream which is not synchronized before the event is recorded on wgrad_stream. This could cause a data race if the two features are combined.
| """Delay the weight gradient computation for TE Grouped GEMM MoE experts. | ||
| When enabled with FSDP, the expert weight gradients are computed on a separate | ||
| CUDA stream after the data gradients finish, allowing overlap of wgrad compute | ||
| with the backward pass of earlier layers. The FSDP gradient reduce-scatter for |
There was a problem hiding this comment.
The docstring says the wgrad overlaps "with the backward pass of earlier layers," but looking at the implementation in _RegisterDelayedWgradForExperts.backward, the main backward stream synchronizes (current_stream().wait_event(event)) before returning — so earlier layers' backward cannot start until wgrad finishes. The actual overlap is between the wgrad computation and the A2A combine backward (dispatch backward) within the same layer.
| """Delay the weight gradient computation for TE Grouped GEMM MoE experts. | |
| When enabled with FSDP, the expert weight gradients are computed on a separate | |
| CUDA stream after the data gradients finish, allowing overlap of wgrad compute | |
| with the backward pass of earlier layers. The FSDP gradient reduce-scatter for | |
| """Delay the weight gradient computation for TE Grouped GEMM MoE experts. | |
| When enabled, the expert weight gradients are computed on a separate | |
| CUDA stream after the data gradients finish, allowing overlap of wgrad compute | |
| with the A2A combine communication within the same MoE layer. When used with | |
| FSDP, the gradient reduce-scatter for expert parameters is deferred until the | |
| delayed wgrad computation completes. | |
| This requires transformer_engine with GroupedLinear support (TE >= 2.3.0). |
| @@ -0,0 +1,202 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Nit: copyright year should be 2026.
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. |
1 participant
What does this PR do ?
PR for dev: #3766
Problem
In MoE models, the expert weight gradient (wgrad) computation during backward is serialized on the main CUDA stream. This blocks the data gradient (dgrad) from flowing to earlier layers until the expert wgrad finishes, even though there is no data dependency between them. The result is wasted GPU cycles — earlier layers' backward pass sits idle waiting for expert wgrad to complete.
With FSDP, this is further compounded because the gradient reduce-scatter for expert parameters is also blocked on the same critical path.
Solution
This PR introduces a new flag
--delay-wgrad-compute-for-te-grouped-gemmthat separates the expert wgrad computation from the main backward stream:Two autograd functions are inserted into the MoE layer's forward graph:
_RecordExpertDgradCompletion— placed before the expert computation; during backward, it records a CUDA event once the expert dgrad is done._RegisterDelayedWgradForExperts— placed at the dispatch boundary; during backward, it waits on the dgrad event, then launchesbackward_dw()on a dedicated CUDA stream, and synchronizes back to the main stream before proceeding.FSDP integration — When used with MegatronFSDP, expert parameters are marked with
_fsdp_delay_grad_reduce = Trueso the normal post-accumulate-grad hook skips them. A callback is registered viaregister_process_expert_grads_fn()that triggers the FSDP reduce-scatter for expert parameters only after the delayed wgrad computation completes.TE GroupedLinear is configured with
delay_wgrad_compute=True, which tells Transformer Engine to skip wgrad during the normal autograd backward and instead wait for an explicitbackward_dw()call.How to enable
Requirements:
moe_grouped_gemmenabled (not legacy grouped gemm)--delay-wgrad-compute(the existing A2A-overlap-based delay)--overlap-moe-expert-parallel-commWorks with both FSDP and 3-D parallelism (TP/EP/PP).
What is achieved
The expert wgrad computation runs on a separate CUDA stream, overlapping with the EP communication within the same transformer layer. This reduces the wall-clock time of the backward pass without changing numerical results — the feature is bit-exact with the non-delayed baseline (verified by unit tests comparing per-step losses and final weights over multiple optimizer steps).
Changes
megatron/core/model_parallel_config.pydelay_wgrad_compute_for_te_grouped_gemmmegatron/core/transformer/transformer_config.pymegatron/core/transformer/moe/moe_layer.pyregister_process_expert_grads_fncallbackmegatron/core/extensions/transformer_engine.pydelay_wgrad_compute=Trueto TE GroupedLinear when the new flag is setmegatron/core/distributed/fsdp/.../megatron_fsdp.pytests/unit_tests/a2a_overlap/test_delay_wgrad_compute.pyTest plan
test_delay_wgrad_compute_for_te_grouped_gemm— full-model training loop (forward → backward → optimizer) comparing delayed vs. non-delayed acrossnum_layers × shared_experts × dispatcher_type × fp8_flagtest_delay_wgrad_compute_for_te_grouped_gemm_with_fsdp— same comparison with MegatronFSDP wrapping (fully_shard_model+fully_shard_optimizer), verifying the deferred reduce-scatter pathContribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.